import cv2
import os
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import RPi.GPIO as GPIO
import time

class FaceCapture:
    # 添加类变量来记录是否是第一次采集
    first_capture = True

    def __init__(self, user_name, user_role, sample_count=15):
        # 初始化GPIO
        GPIO.setmode(GPIO.BCM)
        self.pin_map = {18: {'key': 32, 'last_state': GPIO.HIGH}}
        for pin in self.pin_map:
            GPIO.setup(pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)
        
        # 用户信息
        self.user_name = user_name
        self.user_role = user_role
        self.sample_count = sample_count
        self.data_dir = 'data'
        
        # 只在第一次采集时清空整个文件夹
        if FaceCapture.first_capture:
            if os.path.exists(self.data_dir):
                # 使用shutil.rmtree彻底删除整个文件夹
                import shutil
                shutil.rmtree(self.data_dir)
            FaceCapture.first_capture = False
        
        # 创建数据文件夹（如果不存在）
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
        # 只清空当前用户的数据
        user_dir = os.path.join(self.data_dir, f'{self.user_name}_{self.user_role}')
        if os.path.exists(user_dir):
            for root, dirs, files in os.walk(user_dir, topdown=False):
                for file in files:
                    os.remove(os.path.join(root, file))
                for dir in dirs:
                    os.rmdir(os.path.join(root, dir))
            os.rmdir(user_dir)
        
        # 初始化摄像头
        self.cap = cv2.VideoCapture(0)
        self.font = ImageFont.truetype('wenjian/simhei.ttf', 20)
        cv2.namedWindow('Face Capture', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Face Capture', 800, 480)
        cv2.moveWindow('Face Capture', 0, 0)
        
        # 加载人脸检测模型
        self.face_net = cv2.dnn.readNetFromCaffe(
            "wenjian/deploy.prototxt",
            "wenjian/res10_300x300_ssd_iter_140000.caffemodel"
        )

    def detect_faces_dnn(self, frame):
        """检测人脸并返回坐标"""
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                    (300, 300), (104.0, 177.0, 123.0))
        self.face_net.setInput(blob)
        detections = self.face_net.forward()
        faces = []
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                faces.append((startX, startY, endX - startX, endY - startY))
        return faces

    def check_keypress(self):
        """检测GPIO按键"""
        for pin, config in self.pin_map.items():
            current_state = GPIO.input(pin)
            if current_state != config['last_state']:
                time.sleep(0.02)  # 20ms防抖
                config['last_state'] = current_state
                if current_state == GPIO.HIGH:
                    return config['key']
        return None

    def capture_faces(self):
        """主采集循环"""
        user_dir = os.path.join(self.data_dir, f'{self.user_name}_{self.user_role}')
        if not os.path.exists(user_dir):
            os.makedirs(user_dir)
        
        count = 0
        print(f"开始采集 {self.user_name} 的人脸数据...")
        
        while count < self.sample_count:
            ret, frame = self.cap.read()
            if not ret:
                continue
            
            frame = cv2.rotate(frame, cv2.ROTATE_180)
            frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(frame_pil)
            
            faces = self.detect_faces_dnn(frame)
            
            if len(faces) == 1:
                x, y, w, h = faces[0]
                draw.rectangle([x, y, x+w, y+h], outline=(0, 255, 0), width=2)
                draw.text((x, y-25), f"检测到人脸，按空格拍摄 ({count}/{self.sample_count})",
                         font=self.font, fill=(0, 255, 0))
            else:
                draw.text((10, 20), "未检测到人脸", font=self.font, fill=(255, 0, 0))
            
            draw.text((10, 40), f"当前用户: {self.user_name} ({self.user_role})",
                     font=self.font, fill=(0, 255, 0))
            draw.text((10, 60), "ESC退出 | 空格拍摄", 
                     font=self.font, fill=(0, 255, 0))
            
            cv2.imshow('Face Capture', cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR))
            
            key = cv2.waitKey(30)
            gpio_key = self.check_keypress()
            if gpio_key == 32:  # GPIO 18触发
                key = 32
                
            if key == 32 and len(faces) == 1:  # 空格键拍照
                x, y, w, h = faces[0]
                face_roi = frame[y:y+h, x:x+w]
                gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
                cv2.imwrite(os.path.join(user_dir, f'{self.user_name}_{self.user_role}_{count}.jpg'), gray)
                count += 1
                print(f"已保存第 {count} 张照片")
            elif key == 27:  # ESC退出
                break
        
        print(f"{self.user_name} 的人脸数据采集完成!")
        self.cleanup()

    def cleanup(self):
        """释放资源"""
        self.cap.release()
        cv2.destroyAllWindows()
        GPIO.cleanup()

if __name__ == "__main__":
    # 示例用法
    capture = FaceCapture(
        user_name="张三", 
        user_role="非遗传承人",
        sample_count=15
    )
    capture.capture_faces()
    capture = FaceCapture(
        user_name="李四", 
        user_role="小狗蛋",
        sample_count=5
    )
    capture.capture_faces()
